import socket
import time
import struct
import random
from threading import Timer
from threading import Lock
from collections import deque

NIL = 0
ACK = 1
NAK = 2

VER = "GBN"


class DataStream:
    def __init__(self, data: bytes, pktsize: int) -> None:
        self.data = data
        self.length = len(data)
        self.readptr = 0
        self.pktsize = pktsize
        self.amount = self.length // self.pktsize + 1
        pass

    def __iter__(self):
        return self

    def __next__(self):
        if self.readptr < self.length:
            nxtptr = self.readptr + min(self.pktsize, self.length - self.readptr)
            tmpdata = self.data[self.readptr : nxtptr]
            self.readptr = nxtptr
            return tmpdata
        elif self.readptr == self.length:
            self.readptr += 1
            print("void sent")
            return b""
        else:
            print("data finished")
            raise StopIteration

    pass


class GBNClient:
    def __init__(
        self,
        filename: str,
        pktsize: int,
        target: tuple[str, int],
        errorsize: int,
        timeout: int,
        losssize: int,
        maxwwidth: int,
    ) -> None:
        self.datastream = DataStream(open(filename, "rb").read(), pktsize)
        self.pktsize = pktsize
        self.target = target
        self.clientSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 0 -> waiting for call 0 from above
        # 1 -> waiting for ack/nak 0
        # 2 -> waiting for call 1 from above
        # 3 -> waiting for ack/nak 1
        self.state = 0
        self.errorsize = errorsize

        self.timeout = timeout
        self.timer = Timer(interval=self.timeout, function=self.timercallback)
        self.losssize = losssize

        self.maxwwidth = maxwwidth
        self.window = deque(maxlen=self.maxwwidth)
        self.windowlock = Lock()
        pass

    def timercallback(self):
        # 重发所有未被确认的包
        with self.windowlock:
            for seq, data in self.window:
                self.sendpkt(self.biterror(self.makepkt(data, NIL, seq)))
        self.timer = Timer(interval=self.timeout, function=self.timercallback)
        self.timer.start()
        pass

    def makepkt(self, rawdata: bytes, confirmation: int, sequence: int):
        # gbn 报文格式:
        # ack/nak 1byte
        # sequence 2byte
        # length 2byte
        # checksum 2byte
        pkt = struct.pack(
            "!BHHH%ds" % len(rawdata), confirmation, sequence, len(rawdata), 0, rawdata
        )
        csum = self.checksum(pkt)
        return struct.pack(
            "!BHHH%ds" % len(rawdata),
            confirmation,
            sequence,
            len(rawdata),
            csum,
            rawdata,
        )

    def sendpkt(self, pkt: bytes):
        if random.randint(1, self.losssize) == 1:
            print("packet loss at ", self.datastream.readptr)
            pass
        else:
            self.clientSocket.sendto(self.biterror(pkt), self.target)

    def checksum(self, data: bytes) -> int:
        csum = 0
        for i in range(1, len(data), 2):
            csum += int.from_bytes(data[i - 1 : i + 1], "big")
        if len(data) % 2 == 1:
            csum += data[-1] << 8 & 0xFF00
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = ~csum & 0xFFFF
        # return csum
        return csum >> 8 | ((csum << 8) & 0xFF00)

    def extract(self, pkt: bytes) -> tuple[int, int, int, int, bytes]:
        confirmation, sequence, length, csum = struct.unpack("!BHHH", pkt[:7])
        data = struct.unpack_from("!%ds" % length, pkt, 7)[0]
        return confirmation, sequence, length, csum, data

    def biterror(self, pkt: bytes):
        if random.randint(1, self.errorsize) == 1:
            pos = random.randint(0, len(pkt) - 1)
            off = random.randint(0, 7)
            npkt = bytearray(pkt)
            npkt[pos] ^= 1 << off
            return bytes(npkt)
        return pkt

    def valid(self, pkt: bytes) -> bool:
        return self.checksum(pkt) == 0

    def run(self):
        dataiter = iter(self.datastream)
        partdata = None

        baseseqnum = 1
        nextseqnum = 1
        datafinished = False
        checkfinished = False
        # while checkfinished != True:
        # try:
        # send part
        while datafinished != True or checkfinished != True:
            with self.windowlock:
                for i in range(self.window.maxlen - len(self.window)):
                    try:
                        partdata = next(dataiter)
                        print(
                            f"data with seq {nextseqnum} send. current window: [{baseseqnum}, {nextseqnum}]"
                        )
                        self.sendpkt(
                            self.biterror(self.makepkt(partdata, NIL, nextseqnum))
                        )
                        self.window.append((nextseqnum, partdata))
                        if baseseqnum == nextseqnum:
                            # set timer
                            self.timer = Timer(
                                interval=self.timeout, function=self.timercallback
                            )
                            self.timer.start()
                            pass
                        nextseqnum += 1
                    except StopIteration:
                        datafinished = True
                        break

            # rcv part
            pktrcvd = self.clientSocket.recvfrom(1024)[0]
            if self.valid(pktrcvd):
                # ack.
                confirmation, sequence, length, csum, data = self.extract(pktrcvd)

                if sequence < baseseqnum:
                    # 特殊情况；第一个包出错，返回ack的seq小于baseseqnum
                    pass
                else:
                    baseseqnum = sequence + 1
                    with self.windowlock:
                        while len(self.window) and self.window.popleft()[0] != sequence:
                            pass

                print(f"data accepted to seq {sequence}.")

                with self.windowlock:
                    if datafinished and len(self.window) == 0:
                        print(f"all ack received.")
                        checkfinished = True

                self.timer.cancel()
                if baseseqnum == nextseqnum:
                    # 所有已经发出去的包都受到了确认
                    # 停止计时器并不再重新启动
                    pass
                else:
                    # 重新启动计时器
                    self.timer = Timer(
                        interval=self.timeout, function=self.timercallback
                    )
                    self.timer.start()
                    pass
            else:
                pass


c = GBNClient(
    filename="komeiji.png",
    pktsize=960,
    target=("127.0.0.1", 65432),
    errorsize=100,
    timeout=0.1,
    losssize=100,
    maxwwidth=5,
)
c.run()
